import os
import aiomysql
from typing import Optional
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()


class DatabaseConnection:
    """数据库连接管理器"""

    def __init__(self):
        self.host = os.getenv("DB_HOST")
        self.port = int(os.getenv("DB_PORT"))
        self.user = os.getenv("DB_USER")
        self.password = os.getenv("DB_PASSWORD")
        self.database = os.getenv("DB_NAME")
        self.minsize = 1
        self.maxsize = 10
        self._pool: Optional[aiomysql.Pool] = None

    async def create_pool(self):
        """创建数据库连接池"""
        if self._pool is None:
            self._pool = await aiomysql.create_pool(
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                db=self.database,
                minsize=self.minsize,
                maxsize=self.maxsize,
                charset="utf8mb4",
                autocommit=True,
            )
        return self._pool

    async def get_pool(self):
        """获取连接池"""
        if self._pool is None:
            await self.create_pool()
        return self._pool

    async def close_pool(self):
        """关闭连接池"""
        if self._pool:
            self._pool.close()
            await self._pool.wait_closed()
            self._pool = None
